import torch


class LinearRegression1:
    """
    训练一个线性回归模型
    输入、权重、偏置项和输出都是通过模拟
    输入是100个维度为3的数据，每个元素都是0-1之间
    对应三个维度的权重是1.1，2.2，3.3
    偏置项是4.4
    输出：WX+B+一些随机数

    目标：通过训练还原权重和偏置项，训练前，我们知道的前置信息是输入和输出，且线性模型能够很好拟合输入和输出的关系
    """
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self._init_data()

    def _init_data(self):
        """
        创建数据，input具有3个特征共一百条
        :return:
        """
        self.inputs = torch.rand(100, 3) # 创建一个shape为(100,3)的tensor，里边每个元素都是0-1之间
        self.weights = torch.tensor([[1.1], [2.2], [3.3]]) # 初始化权重
        self.bias = torch.tensor(4.4) # 创建一个偏置项
        self.targets = self.inputs @ self.weights + self.bias + 0.1 * torch.randn(100, 1) # 创建一个线性回归的模型，并加上一些噪声（ 0.1 * torch.randn(100, 1)）

    def run(self):
        """
        开始训练
        :return:
        """
        w = torch.rand(3, 1, requires_grad=True, device=self.device) # 创建初始权重，并启用梯度追踪
        b = torch.rand(1, requires_grad=True, device=self.device) # 创建初始偏置项，并启用梯度追踪

        inputs = self.inputs.to(self.device) # 将数据移至相同设备
        targets = self.targets.to(self.device) # 将数据移至相同设备

        # 设置超参数
        epoch = 10000 # 迭代次数
        lr = 0.003  # 学习率

        for i in range(epoch):
            outputs = inputs @ w + b # 计算100个输入的预测输出
            loss = torch.mean(torch.square(outputs - targets)) # 计算预测输出和真实输出的均方误差，损失函数，值是个标量
            print("loss:", loss.item())

            loss.backward() # 对 损失函数求梯度

            with torch.no_grad():
                w -= lr * w.grad
                b -= lr * b.grad

            w.grad.zero_()
            b.grad.zero_()

        print("训练后的权重 w:", w)
        print("训练后的偏置 b:", b)

linear_regression = LinearRegression1()
linear_regression.run()